#define PI 3.141592653589
#include "dwf_4d_invert.h"

/**
 * Fourier Transform a domain-wall fermion
 *
 * A fourier transform function that will transform the 's' coordinate of a domain-wall fermion into
 * 'k' space.
 *
 * @param eta Source fermion to be transformed
 * @param fteta Reference to transformed fermion
 *
 */
void dwf_4d_invert::fourierTransform(multi1d<LatticeFermion> &eta, multi1d<LatticeFermion> &fteta)
{

    int Ls = eta.size();
    // Loop elements in FT'd fermion
    for ( int k = 0; k < Ls; k++ )
    {
        fteta[k] = zero;
        // Just the two edge terms included
        for ( int s = 0; s < Ls; s+=(Ls-1) )
        {
            double arg = 2.0*PI*(k+0.5)*s/(double)Ls;
            Complex exp = cmplx((Real)std::cos(arg),(Real)std::sin(arg));
            fteta[k] += exp*eta[s];
        }
    }
}



/**
 * Inverse Fourier Transform a domain-wall fermion
 *
 * Inverse fourier transform a domain-wall fermion from 'k' space back to
 * 's' space of length Ls
 *
 * @param psi Source fermion in 'k' space
 * @param iftpsi Reference to resulting fermion in 's' space
 *
 */
void dwf_4d_invert::fourierTransformInv(multi1d<LatticeFermion> &psi, multi1d<LatticeFermion> &iftpsi)
{

    int Ls = psi.size();
    // Loop elements in Inverse FT'd fermion
    for ( int s = 0; s < Ls; s++ )
    {
        iftpsi[s] = zero;
        for ( int k = 0; k < Ls; k++ )
        {
            double arg = -2.0*PI*(k+0.5)*s/(double)Ls;
            Complex exp = cmplx((Real)std::cos(arg),(Real)std::sin(arg));
            iftpsi[s] += (exp*psi[k])/Ls;
        }
    }
}


/**
 * Perform a complete inversion of Ddwf 5d operator with fermions existing in real space
 *
 * @param eta Source fermion
 * @param psi Solution Fermion
 * @param u Gauge field
 * @param M5 Mass M5
 * @param device Device number of GPU on which the inversion should be performed
 */

void dwf_4d_invert::fullInversion(multi1d<LatticeFermion> &src, multi1d<LatticeFermion> &sol, multi1d<LatticeColorMatrix> &u, double M5, int dagger)
{
    multi1d<LatticeFermion> tmp(src.size());

    dwf_4d_invert::fourierTransform(src, sol);
    dwf_4d_invert::domainWallFermionInvert(sol, tmp, u, M5, dagger);
    dwf_4d_invert::fourierTransformInv(tmp, sol);
}


/**
 * Perform the inversion of \f$ D_{dwf}\psi = \eta \f$ to find \f$ \psi \f$ where \f$ \psi \f$ and \f$ \eta \f$ are in Fourier space
 *
 * @param eta Source fermion
 * @param psi Solution Fermion
 * @param u Gauge field
 * @param M5 Mass M5
 * @param device Device number of GPU on which the inversion should be performed
 */
void dwf_4d_invert::domainWallFermionInvert(multi1d<LatticeFermion> &eta, multi1d<LatticeFermion> &psi, multi1d<LatticeColorMatrix> &u,double M5, int dagger){

    ////////////////////////////////////////////////
    // QUDA PARAMETERS AND SETUP
    ////////////////////////////////////////////////

    // TODO: Better way of finding QDP precision
    QudaPrecision cpu_prec;
    switch ( sizeof(Real(0)) )
    {
        case 4:
            cpu_prec = QUDA_SINGLE_PRECISION;
            break;
        case 8:
            cpu_prec = QUDA_DOUBLE_PRECISION;
            break;
        default:
            std::cerr << "Unable to determine QDP++ precision.";
            exit(1);
    }

    // Using GTX 480, single precision
    QudaPrecision cuda_prec = QUDA_SINGLE_PRECISION;
    QudaPrecision cuda_prec_sloppy = QUDA_SINGLE_PRECISION;

    QudaGaugeParam gauge_param = newQudaGaugeParam();
    QudaInvertParam inv_param = newQudaInvertParam();

    multi1d<int> Lat = Layout::lattSize();
    gauge_param.X[0] = Lat[0];
    gauge_param.X[1] = Lat[1];
    gauge_param.X[2] = Lat[2];
    gauge_param.X[3] = Lat[3];

    gauge_param.anisotropy = 1.0;
    gauge_param.type = QUDA_WILSON_LINKS;
    gauge_param.gauge_order = QUDA_QDP_GAUGE_ORDER;
    gauge_param.t_boundary = QUDA_ANTI_PERIODIC_T;

    gauge_param.cpu_prec = cpu_prec;
    gauge_param.cuda_prec = cuda_prec;
    gauge_param.reconstruct = QUDA_RECONSTRUCT_NO;
    gauge_param.cuda_prec_sloppy = cuda_prec_sloppy;
    gauge_param.reconstruct_sloppy = QUDA_RECONSTRUCT_NO;
    gauge_param.gauge_fix = QUDA_GAUGE_FIXED_NO;

    inv_param.dslash_type = QUDA_TWISTED_MASS_DSLASH;
    inv_param.inv_type = QUDA_CG_INVERTER;
    inv_param.twist_flavor = QUDA_TWIST_PLUS;

    inv_param.tol = 1e-6;
    inv_param.maxiter = 10000;
    inv_param.reliable_delta = 1e-2;

    inv_param.solution_type = QUDA_MAT_SOLUTION;
    inv_param.solve_type = QUDA_DIRECT_SOLVE;
    inv_param.matpc_type = QUDA_MATPC_EVEN_EVEN;
    if ( dagger )
        inv_param.dagger = QUDA_DAG_YES;
    else
        inv_param.dagger = QUDA_DAG_NO;

    inv_param.mass_normalization = QUDA_MASS_NORMALIZATION;

    inv_param.cpu_prec = cpu_prec;
    inv_param.cuda_prec = cuda_prec;
    inv_param.cuda_prec_sloppy = cuda_prec_sloppy;
    inv_param.preserve_source = QUDA_PRESERVE_SOURCE_YES;
    inv_param.dirac_order = QUDA_DIRAC_ORDER;

    gauge_param.ga_pad = 0; // 24*24*24;
    inv_param.sp_pad = 0;   // 24*24*24;
    inv_param.cl_pad = 0;   // 24*24*24;

    inv_param.verbosity = QUDA_VERBOSE;

    ///////////////////////////////////////
    // End config and begin preparing data
    ///////////////////////////////////////
    void* gauge[4];
    for(int mu=0; mu < Nd; mu++)
    { 
        // Pointer to first element of Gauge field in each direction
        gauge[mu] = (void *)&(u[mu].elem(all.start()).elem().elem(0,0).real());
    }

    // initialize the QUDA library and import gauge field
    loadGaugeQuda(gauge, &gauge_param);


    //////////////////////////////////////////////////////////
    // Performing the inversion for each 'k'
    //////////////////////////////////////////////////////////

    int Ls = eta.size();
    for ( int k = 0; k < Ls; k++ )
    {
        double alpha_k = 2.0*PI*(k+0.5)/(double)Ls;
        double mass = -M5 + 1 - std::cos(alpha_k);
        //double ATerm = -M5 + 5 - std::cos(alpha_k);
        inv_param.kappa = (1.0 / (2.0*(1.0 + 3.0/gauge_param.anisotropy + mass)));
        inv_param.mu = -std::sin(alpha_k);


        // Again, pointer to first element of each 'k'
        void* spinorIn = (void *)&(eta[k].elem(all.start()).elem(0).elem(0).real());
        void* spinorOut = (void *)&((psi[k]).elem(all.start()).elem(0).elem(0).real());
        // Perform the inversion
        invertQuda(spinorOut,spinorIn, &inv_param);
        if ( inv_param.verbosity == QUDA_VERBOSE )
            printf("Device memory used:\n   Spinor: %f GiB\n    Gauge: %f GiB\n", 
                    inv_param.spinorGiB, gauge_param.gaugeGiB);
    }

}

